# gomoku_ai/game/gomoku_game.py
import numpy as np
import pygame
import sys


class GomokuGame:
    def __init__(self, board_size=9, visual_mode=False):
        self.board_size = board_size
        self.visual_mode = visual_mode
        self.visual_board = None

        if self.visual_mode:
            self.visual_board = self._init_visual_board()

    def _init_visual_board(self):
        """初始化图形界面"""
        from utils.visual_board import VisualBoard
        return VisualBoard(self.board_size)

    def get_init_board(self):
        """获取初始棋盘"""
        return np.zeros((self.board_size, self.board_size), dtype=int)

    def get_board_size(self):
        """获取棋盘大小"""
        return (self.board_size, self.board_size)

    def get_action_size(self):
        """获取动作空间大小"""
        return self.board_size * self.board_size

    def get_next_state(self, board, player, action):
        """执行动作，返回下一个状态"""
        b = np.copy(board)
        b[action // self.board_size][action % self.board_size] = player
        return b, -player

    def get_valid_moves(self, board):
        """获取所有有效的动作"""
        valids = np.zeros(self.get_action_size())
        valids[board.reshape(-1) == 0] = 1
        return valids

    def get_game_ended(self, board):
        """检查游戏是否结束，返回赢家 (1: 黑棋, -1: 白棋, 0: 未结束, 0.1: 平局)"""
        # 检查横向、纵向、对角线是否有五子连珠
        for i in range(self.board_size):
            for j in range(self.board_size - 4):
                if board[i, j] != 0 and all(board[i, j + k] == board[i, j] for k in range(5)):
                    return board[i, j]
                if board[j, i] != 0 and all(board[j + k, i] == board[j, i] for k in range(5)):
                    return board[j, i]

        # 检查对角线
        for i in range(self.board_size - 4):
            for j in range(self.board_size - 4):
                if board[i, j] != 0 and all(board[i + k, j + k] == board[i, j] for k in range(5)):
                    return board[i, j]
                if board[i + 4, j] != 0 and all(board[i + 4 - k, j + k] == board[i + 4, j] for k in range(5)):
                    return board[i + 4, j]

        # 检查是否平局
        if np.count_nonzero(board) == self.board_size * self.board_size:
            return 0.1  # 平局

        return 0  # 游戏未结束

    def get_canonical_form(self, board, player):
        """获取标准形式的棋盘（从当前玩家视角）"""
        return player * board

    def get_symmetries(self, board, pi):
        """获取棋盘和策略的对称形式（旋转和翻转）"""
        assert (len(pi) == self.board_size ** 2)
        pi_board = np.reshape(pi, (self.board_size, self.board_size))
        l = []

        for i in range(1, 5):
            for j in [True, False]:
                newB = np.rot90(board, i)
                newPi = np.rot90(pi_board, i)
                if j:
                    newB = np.fliplr(newB)
                    newPi = np.fliplr(newPi)
                l += [(newB, list(newPi.ravel()))]
        return l

    def string_representation(self, board):
        """棋盘的字符串表示（用于MCTS哈希）"""
        # --- 关键修改：使用更稳健的元组表示法 ---
        # 这确保了不同的棋盘布局总是有不同的哈希值，防止MCTS陷入循环
        return str(board.tolist())

    def render(self, board, game_result=0, current_player=1, info_text=""):
        """渲染棋盘"""
        if self.visual_mode:
            self.visual_board.update(board, game_result, current_player, info_text)
        else:
            self._render_console(board, current_player, info_text)

    def _render_console(self, board, current_player, info_text):
        """在终端渲染棋盘"""
        print("\n" + info_text)
        print("   " + " ".join([str(i).rjust(2) for i in range(self.board_size)]))
        for i in range(self.board_size):
            print(str(i).rjust(2) + " " + " ".join(['. ' if p == 0 else 'X ' if p == 1 else 'O ' for p in board[i]]))
        print(f"当前玩家: {'黑棋(X)' if current_player == 1 else '白棋(O)'}")

    def get_human_move(self, board):
        """获取人类玩家的输入"""
        if not self.visual_mode:
            while True:
                try:
                    move = input("请输入你的走法 (格式: x,y): ")
                    x, y = map(int, move.split(','))
                    if 0 <= x < self.board_size and 0 <= y < self.board_size and board[x, y] == 0:
                        return x * self.board_size + y
                    else:
                        print("无效的走法，请重试。")
                except Exception:
                    print("输入格式错误，请重试。")
        else:
            return self.visual_board.get_human_move(board)

    def wait_for_continue(self, board, game_result):
        """等待用户继续（用于训练模式）"""
        if self.visual_mode:
            # --- 关键修改：提供更明确的提示和更可靠的等待机制 ---
            font = pygame.font.SysFont('simhei', 24)

            # 创建半透明覆盖层
            overlay = pygame.Surface((self.visual_board.screen.get_width(), self.visual_board.screen.get_height()))
            overlay.set_alpha(200)
            overlay.fill((0, 0, 0))
            self.visual_board.screen.blit(overlay, (0, 0))

            # 显示提示信息
            if game_result == 1:
                text = "黑棋获胜！"
            elif game_result == -1:
                text = "白棋获胜！"
            else:
                text = "平局！"

            prompt_text = "点击鼠标或按任意键继续下一局..."

            text_surface = font.render(text, True, (255, 255, 255))
            prompt_surface = font.render(prompt_text, True, (200, 200, 200))

            text_rect = text_surface.get_rect(
                center=(self.visual_board.screen.get_width() // 2, self.visual_board.screen.get_height() // 2 - 20))
            prompt_rect = prompt_surface.get_rect(
                center=(self.visual_board.screen.get_width() // 2, self.visual_board.screen.get_height() // 2 + 20))

            self.visual_board.screen.blit(text_surface, text_rect)
            self.visual_board.screen.blit(prompt_surface, prompt_rect)

            pygame.display.flip()

            # 等待用户交互
            waiting = True
            while waiting:
                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        pygame.quit()
                        sys.exit()
                    if event.type == pygame.MOUSEBUTTONDOWN or event.type == pygame.KEYDOWN:
                        waiting = False
        else:
            input("按回车键继续下一局...")

    def close(self):
        """关闭图形界面"""
        if self.visual_board:
            self.visual_board.close()
